import sqlite3
import os


class Database:
    def __init__(self):
        self.conn = sqlite3.connect(r'../action.db')
        self.cuis = self.conn.cursor()
        self.id_to_class = dict()
        self.imagePath = r'/home/zp1/data/ImageData/'

    def init_id_to_class(self):
        # 设置字典
        with open('../data/my_id_class5.txt', 'r') as f:
            for data in f.readlines():
                l_temp = data.split(':')
                self.id_to_class[l_temp[0]] = l_temp[1]

    def create_action_table(self):
        self.cuis.execute(
            '''
            CREATE TABLE action(
            id TEXT PRIMARY KEY ,
            address TEXT,   
            class TEXT
            '''
        )

    def create_class_table(self):
        pass

    def insert(self):
        # 插入测试
        movie_name_list = sorted(os.listdir(self.imagePath))
        for i in movie_name_list:
            if i[:5] not in self.id_to_class.keys():
                print(i)
                continue
            id = i[:-4]
            address = i
            command = f'insert into action (id,address,class) values ("{id}","{self.imagePath + address}","{self.id_to_class[id]}")'
            print(command)
            self.cuis.execute(command)
        self.conn.commit()
        print('数据插入成功')

    def read(self):
        cursor = self.conn.execute("select * from action")
        for row in cursor:
            print('ID = ', row[0], '  address = ', row[1], '  class = ', row[2])
        print('Operation done successfully')

    def __del__(self):
        self.conn.close()


if __name__ == '__main__':
    database = Database()
    # database.create_action_table()
    database.init_id_to_class()
    database.insert()
    database.read()
